# !/usr/bin/env python
# -*- coding: utf-8 -*-
from pyspark.sql import SparkSession
from pyspark.sql.functions import isnull, when, count, col
from pyspark.ml.feature import VectorAssembler, StringIndexer
from pyspark.ml.classification import RandomForestClassifier, DecisionTreeClassifier, GBTClassifier
from pyspark.ml.evaluation import MulticlassClassificationEvaluator
# 请确保你的 DataFrame 包含一个名为 Label 的列, 这是 Species(品种) 的列。
# 如果 label 是字符串类型的分类特征, 你可能需要使用 StringIndexer 进行索引。
from pyspark.sql.types import DoubleType, StringType

'''
任务类型: 随机森林回归主要用于回归任务。在回归任务中, 算法试图预测一个连续的数值输出, 而不是一个离散的类别。
输出: 随机森林回归的输出是一个连续的数值, 表示输入数据的预测结果。
算法原理: 随机森林回归同样基于决策树, 但在回归任务中, 每个决策树的输出是一个实数值。最终的预测结果是多个决策树输出的平均值或加权平均值。

字段说明: SepalLength(花萼长度), SepalWidth(花萼宽度), PetalLength(花瓣长度), PetalWidth(花瓣宽度), Species(品种).

品种说明: Setosa(山鸢尾), Versicolor(变色鸢尾), Virginical(维吉尼亚鸢尾).
'''
'''
在分类(classification)问题的模型评估中，常用的评测指标有以下7个：
准确率(accuracy)
精确率(precision)
召回率(recall)
F1-Score
ROC曲线
P-R曲线
AUC面积
https://blog.csdn.net/michael_f2008/article/details/104938824
'''


def get_classifier(name, labelCol, featuresCol):
    if name == 'RandomForest':
        # 随机森林分类器
        return RandomForestClassifier(labelCol=labelCol, featuresCol=featuresCol, maxDepth=5)
    elif name == 'DecisionTree':
        # 决策树分类器
        return DecisionTreeClassifier(labelCol=labelCol, featuresCol=featuresCol, maxDepth=5)
    elif name == 'GBT':
        # 梯度提升树分类器模型
        return GBTClassifier(labelCol=labelCol, featuresCol=featuresCol, maxDepth=5)
    return None


if __name__ == '__main__':
    spark = SparkSession.builder.appName("classifier_learn").master("local[*]").getOrCreate()

    # 加载数据
    df = spark.read.option("header", "true").option("encoding", "utf-8").csv("../data/iris.csv")
    df.printSchema()

    # 缺失数据处理
    df.select([count(when(isnull(c), c)).alias(c) for c in df.columns]).show()

    # 转换 Spark 中 DateFrame 数据类型。
    df = df.withColumn("SepalLength", col("SepalLength").cast(DoubleType()))
    df = df.withColumn("SepalWidth", col("SepalWidth").cast(DoubleType()))
    df = df.withColumn("PetalLength", col("PetalLength").cast(DoubleType()))
    df = df.withColumn("PetalWidth", col("PetalWidth").cast(DoubleType()))
    df = df.withColumn("Species", col("Species").cast(StringType()))

    # 输出 Spark 中 DataFrame 字段和数据类型。
    print("Changed Spark DataFrame Data Type:")
    df.printSchema()

    # 使用 StringIndexer 转换 Species 列。
    indexer = StringIndexer(inputCol="Species", outputCol="SpeciesIdx")
    # 拟合并转换数据。
    indexDF = indexer.fit(df).transform(df)

    # 输出 StringIndexer 的转换效果。
    print("The Effect of StringIndexer:")
    indexDF.show()

    # 将 SpeciesIdx 列复制为 Label 列。
    dataset = indexDF.withColumn("Label", col("SpeciesIdx"))

    # VectorAssembler 将多个特征合并为一个特征向量。
    required_features: list = ["SepalLength", "SepalWidth", "PetalLength", "PetalWidth"]
    assembler = VectorAssembler(inputCols=required_features, outputCol="Features")

    transformed_data = assembler.transform(dataset)
    transformed_data.show()

    # 将数据随机分成训练集和测试集，并设置可重复性的种子。
    (training_data, test_data) = transformed_data.randomSplit([0.8, 0.2], seed=2020)
    print("训练数据集总数: " + str(training_data.count()))
    print("测试数据集总数: " + str(test_data.count()))
    # GBT只适用于二分类
    # class_arr = ["DecisionTree", "RandomForest", "GBT"]
    class_arr = ["DecisionTree", "RandomForest"]
    for name in class_arr:
        print(f"==================={name} classifier=====================")
        rf = get_classifier(name, "Label", "Features")
        model = rf.fit(training_data)
        rf_predictions = model.transform(test_data)
        rf_predictions.show()
        #  评估随机森林分类器模型
        multi_evaluator = MulticlassClassificationEvaluator(labelCol="Label", metricName='accuracy')
        print(name + ' classifier Accuracy:', multi_evaluator.evaluate(rf_predictions))
    spark.stop()
